{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "source": "# Weather Forecast Agent with Pydantic AI & Nebius\n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1ES51c3YKcMSpj6XC3wV89IzbULxusKyy?usp=sharing)\n\nThis notebook demonstrates how to build a simple weather forecast agent using Pydantic AI and Nebius Token Factory. The agent fetches the current weather forecast for a specified city by searching DuckDuckGo with the meta-llama/Meta-Llama-3.1-70B-Instruct model.\n\n[Nebius Token Factory](https://studio.nebius.ai) provides access to many state-of-the-art LLM models. Check out the full list of models here.\n\nVisit https://studio.nebius.ai/ and sign up to get an API key.",
   "metadata": {
    "id": "zwGbdai7pCtV"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Step 1: Install Dependencies"
   ],
   "metadata": {
    "id": "vRmnshrBpzag"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V9TBh9onooMR"
   },
   "outputs": [],
   "source": [
    "!pip install pydantic_ai 'pydantic-ai-slim[duckduckgo]'"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Step 2: Enter Your API Key\n",
    "\n",
    "For security, we’ll prompt you to enter your Nebius API key instead of hardcoding it, since Colab doesn’t natively support .env files."
   ],
   "metadata": {
    "id": "-l5SjhPvp4Av"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# set API key in env or in llm\n",
    "import os\n",
    "\n",
    "os.environ[\"NEBIUS_API_KEY\"] = \"Your Nebius API Key\""
   ],
   "metadata": {
    "id": "1Uf87E7iqBez"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ],
   "metadata": {
    "id": "RtuNoP5IuwYG"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "## Step 3: Import Required Modules\n\nWe import essential modules from pydantic_ai to work with Nebius Token Factory and DuckDuckGo search.\n",
   "metadata": {
    "id": "UsVNoBAIqWI2"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "from pydantic_ai import Agent\n",
    "from pydantic_ai.models.openai import OpenAIModel\n",
    "from pydantic_ai.providers.openai import OpenAIProvider\n",
    "from pydantic_ai.common_tools.duckduckgo import duckduckgo_search_tool"
   ],
   "metadata": {
    "id": "PNQoLK5FqeOB"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Step 4: Initializing the Weather Agent\n",
    "\n",
    "Set up the Nebius AI model and configure the agent to fetch weather forecasts."
   ],
   "metadata": {
    "id": "RZ-63AefrQ3t"
   }
  },
  {
   "cell_type": "code",
   "source": "# Set up the model with the user-provided API key\nmodel = OpenAIModel(\n    model_name='meta-llama/Meta-Llama-3.1-70B-Instruct',\n    provider=OpenAIProvider(\n        base_url='https://api.tokenfactory.nebius.com/v1',\n        api_key=os.environ['NEBIUS_API_KEY']\n    )\n)\n\n# Create the agent with a weather-focused prompt\nagent = Agent(\n    model=model,\n    tools=[duckduckgo_search_tool()],\n    system_prompt=\"You are a weather assistant. Use DuckDuckGo to find the current weather forecast for the requested city.\"\n)",
   "metadata": {
    "id": "3Ko3PpSQraJY"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Step 5: Running a Weather Forecast Query\n",
    "\n",
    "Ask the agent for the weather forecast of a city (e.g., Kolkata). You can change the city to any location you’d like!"
   ],
   "metadata": {
    "id": "WOmmP5VvriS3"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# Define the city\n",
    "city = \"Kolkata\"  # Change this to any city you like!\n",
    "\n",
    "# Run the agent\n",
    "result = agent.run_sync(f\"What is the weather forecast for {city} today?\")\n",
    "\n",
    "# Display the result\n",
    "print(f\"Weather forecast for {city}:\")\n",
    "print(result.data)"
   ],
   "metadata": {
    "id": "r4xp2Rl7rh7d",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "a6b10d0c-4b5c-4c7d-8bf3-0f430731b7ca"
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Weather forecast for Kolkata:\n",
      "The current weather forecast for Kolkata today is mostly sunny with a high of 90°F and a low of 69°F. There is a chance of showers on Wednesday afternoon, but the rest of the week is expected to be mostly clear.\n"
     ]
    }
   ]
  }
 ]
}